#!/usr/bin/env python3
# 指定脚本使用Python 3解释器执行

# 导入必要的库
import rclpy  # ROS 2 Python客户端库，用于创建和管理节点
from rclpy.node import Node  # ROS 2节点基类，所有自定义节点都需继承此类
from livox_interfaces.msg import CustomMsg  # Livox激光雷达的自定义点云消息类型
from sensor_msgs.msg import PointCloud2, PointField  # ROS标准点云消息及字段定义
import numpy as np  # 用于高效的数值计算和数组操作
import struct  # 用于二进制数据的打包与解包，构建PointCloud2消息时使用
import open3d as o3d  # 用于3D点云处理、可视化和聚类分析
from threading import Thread, Lock  # 多线程支持，Lock用于线程间同步
import time  # 用于计时和控制程序执行速度
import multiprocessing  # 用于并行计算，提高点云处理效率

# 定义PointCloud2消息的字段结构，描述点云数据的组成
FIELDS = [
    # X坐标：偏移0字节，32位浮点数，数量1
    PointField(name='x', offset=0, datatype=PointField.FLOAT32, count=1),
    # Y坐标：偏移4字节（紧跟X之后），32位浮点数，数量1
    PointField(name='y', offset=4, datatype=PointField.FLOAT32, count=1),
    # Z坐标：偏移8字节（紧跟Y之后），32位浮点数，数量1
    PointField(name='z', offset=8, datatype=PointField.FLOAT32, count=1),
    # 反射率：偏移12字节（紧跟Z之后），8位无符号整数，数量1
    PointField(name='intensity', offset=12, datatype=PointField.UINT8, count=1),
]
POINT_STEP = 13  # 每个点占用的字节数：3个float(4字节×3) + 1个uint8(1字节) = 13字节
ROW_STEP = 0  # 每行的字节数，后续会根据点云数量动态计算


def process_chunk(chunk, distance_threshold, keep_front, area_half_width, area_half_height, 
                 exclude_area_params):
    """
    并行处理点云数据块的过滤函数，实现距离过滤、区域选择和特定区域屏蔽
    
    参数:
        chunk: 点云数据块 (numpy数组，每行代表一个点[x,y,z,intensity])
        distance_threshold: 距离阈值 (米)，只保留该距离内的点
        keep_front: 是否只保留正前方区域的点
        area_half_width: 前方区域的半宽度 (米)
        area_half_height: 前方区域的半高度 (米)
        exclude_area_params: 要排除的区域参数（包含x_min, x_max, y_min, y_max, z_min, z_max）
        
    返回:
        过滤后的点云数据块
    """
    # 1. 距离过滤：计算每个点到原点的平方距离（避免开平方运算提高效率）
    squared_distances = np.sum(chunk[:, :3] **2, axis=1)
    distance_mask = squared_distances <= (distance_threshold** 2)
    
    # 2. 前方区域过滤：仅保留传感器正前方特定范围内的点
    if keep_front:
        in_front_mask = (chunk[:, 0] > 0)
        in_width_mask = (np.abs(chunk[:, 1]) <= area_half_width)
        in_height_mask = (np.abs(chunk[:, 2]) <= area_half_height)
        
        front_area_mask = np.logical_and(np.logical_and(in_front_mask, in_width_mask), in_height_mask)
        filter_mask = np.logical_and(distance_mask, front_area_mask)
    else:
        filter_mask = distance_mask
    
    # 3. 排除指定区域内的点
    in_exclude_area = (
        (chunk[:, 0] >= exclude_area_params['x_min']) &
        (chunk[:, 0] <= exclude_area_params['x_max']) &
        (chunk[:, 1] >= exclude_area_params['y_min']) &
        (chunk[:, 1] <= exclude_area_params['y_max']) &
        (chunk[:, 2] >= exclude_area_params['z_min']) &
        (chunk[:, 2] <= exclude_area_params['z_max'])
    )
    
    final_mask = np.logical_and(filter_mask, ~in_exclude_area)
    return chunk[final_mask]


class LidarProcessor(Node):
    """LiDAR点云处理器节点类，实现点云数据的接收、处理、发布和可视化"""
    
    def __init__(self):
        """初始化LiDAR处理器节点，设置订阅者、发布者和各种参数"""
        super().__init__('lidar_processor')
        
        # 声明并获取ROS参数：输入点云话题名称
        self.declare_parameter('input_cloud', '/livox/lidar_avia')
        input_topic = self.get_parameter('input_cloud').get_parameter_value().string_value

        # 创建订阅者，订阅Livox激光雷达的自定义点云消息
        self.sub = self.create_subscription(
            CustomMsg, input_topic, self.callback, 10
        )
        
        # 创建发布者，发布处理后的标准PointCloud2格式点云
        self.pcl_pub = self.create_publisher(
            PointCloud2, '/livox/pointcloud2', 10
        )
        
        self.get_logger().info(f"订阅: {input_topic}，发布处理后点云到: /livox/pointcloud2")

        # 初始化Open3D点云对象和可视化窗口
        self.pcd = o3d.geometry.PointCloud()
        self.vis = o3d.visualization.Visualizer()
        self.vis.create_window()
        
        # 设置可视化参数
        self.vis.get_render_option().point_size = 2.0
        self.vis.get_render_option().background_color = np.array([0, 0, 0])
        
        # 线程安全相关
        self.point_cloud_lock = Lock()
        self.is_visualizing = True
        self.is_first_frame = True
        
        # 启动可视化线程
        self.visualization_thread = Thread(target=self.visualization_loop)
        self.visualization_thread.daemon = True
        self.visualization_thread.start()
        
        # 点云计数和保存相关参数
        self.frame_count = 0
        self.save_path = "/path/to/save/pointclouds/"
        self.save_enabled = False
        self.use_fixed_color = True
        
        # 过滤参数
        self.distance_threshold = 50.0
        self.keep_front = True
        
        # 前方区域尺寸参数
        self.area_width = 2
        self.area_height = 2
        self.area_half_width = self.area_width / 2
        self.area_half_height = self.area_height / 2
        
        # 聚类分析参数
        self.cluster_eps = 0.1
        self.cluster_min_points = 30
        self.enable_clustering = True
        self.voxel_size = 0.03  
        
        # 新增：聚类点数可视化范围（只可视化点数在这个范围内的聚类）
        self.cluster_min_visualize = 50    # 最小可视化点数
        self.cluster_max_visualize = 400  # 最大可视化点数
        
        # 并行处理设置
        self.num_processes = max(1, multiprocessing.cpu_count() - 1)
        
        # 排除区域参数
        self.exclude_area_params = {
            'x_min': 0.0,    'x_max': 2.8,
            'y_min': -1,     'y_max': 1.0,
            'z_min': -2,     'z_max': -0.1
        }
        self.get_logger().info(
            f"已启用区域屏蔽: X∈[{self.exclude_area_params['x_min']},{self.exclude_area_params['x_max']}], "
            f"Y∈[{self.exclude_area_params['y_min']},{self.exclude_area_params['y_max']}], "
            f"Z∈[{self.exclude_area_params['z_min']},{self.exclude_area_params['z_max']}]"
        )
        self.get_logger().info(
            f"只可视化点数在 {self.cluster_min_visualize}-{self.cluster_max_visualize} 范围内的聚类"
        )

    def callback(self, msg: CustomMsg):
        """订阅消息的回调函数，处理接收到的LiDAR点云数据"""
        start_time = time.time()
        
        # 转换点云消息为NumPy数组
        points_array = np.zeros((len(msg.points), 4), dtype=np.float32)
        for i, p in enumerate(msg.points):
            points_array[i, 0] = p.x
            points_array[i, 1] = p.y
            points_array[i, 2] = p.z
            points_array[i, 3] = p.reflectivity
        
        # 并行过滤点云
        chunks = np.array_split(points_array, self.num_processes)
        with multiprocessing.Pool(processes=self.num_processes) as pool:
            results = pool.starmap(
                process_chunk, 
                [
                    (chunk, self.distance_threshold, self.keep_front, 
                     self.area_half_width, self.area_half_height, self.exclude_area_params) 
                    for chunk in chunks
                ]
            )
        filtered_points = np.concatenate(results)
        
        # 处理点云并设置颜色
        if len(filtered_points) < 50:
            if self.frame_count % 10 == 0:
                self.get_logger().info(f"过滤后点太少（{len(filtered_points)}点），跳过聚类")
            colors = np.ones((len(filtered_points), 3))  # 白色
        else:
            # 创建Open3D点云对象
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(filtered_points[:, :3])
            # 体素降采样：将空间划分为体素网格，每个体素保留一个点
            downsampled_pcd = pcd.voxel_down_sample(voxel_size=self.voxel_size)
            
            if self.enable_clustering:
                # DBSCAN聚类
                labels = np.array(downsampled_pcd.cluster_dbscan(
                    eps=self.cluster_eps, 
                    min_points=self.cluster_min_points, 
                    print_progress=False
                ))
                
                # 统计每个聚类的点数
                unique_labels, counts = np.unique(labels, return_counts=True)
                cluster_sizes = dict(zip(unique_labels, counts))
                
                # 打印所有聚类的点数
                for label, count in cluster_sizes.items():
                    if label == -1:
                        self.get_logger().info(f"噪声点数量: {count}")
                    else:
                        self.get_logger().info(f"聚类 {label} 的点数: {count}")
                
                # 映射聚类标签到原始点云
                pcd_tree = o3d.geometry.KDTreeFlann(downsampled_pcd)
                full_labels = np.ones(len(filtered_points), dtype=np.int32) * -1
                
                for i in range(len(filtered_points)):
                    [k, idx, _] = pcd_tree.search_knn_vector_3d(filtered_points[i, :3], 1)
                    full_labels[i] = labels[idx[0]]
                
                # 生成颜色（只对符合点数范围的聚类着色）
                max_label = full_labels.max()
                self.get_logger().info(f"检测到 {max_label + 1} 个聚类")
                
                colors = np.zeros((len(filtered_points), 3))  # 默认为黑色（背景色，相当于不可见）
                visible_clusters = 0  # 统计可见聚类数量
                
                for label in range(-1, max_label + 1):
                    # 跳过噪声点
                    if label == -1:
                        continue
                        
                    # 检查当前聚类的点数是否在可视化范围内
                    cluster_size = cluster_sizes.get(label, 0)
                    if self.cluster_min_visualize <= cluster_size <= self.cluster_max_visualize:
                        # 符合条件的聚类：分配颜色
                        label_mask = full_labels == label
                        hue = label / (max_label + 1)
                        rgb = self.hsv_to_rgb(hue, 0.8, 0.9)
                        colors[label_mask] = rgb
                        visible_clusters += 1
                    else:
                        # 不符合条件的聚类：保持黑色（不可见）
                        self.get_logger().info(
                            f"聚类 {label} 点数 {cluster_size} 超出范围，不进行可视化"
                        )
                
                self.get_logger().info(f"可视化聚类数量: {visible_clusters}")
                
            else:
                # 不启用聚类时，基于反射率着色
                intensities = filtered_points[:, 3] / 255.0
                colors = np.zeros((len(filtered_points), 3))
                colors[:, 0] = intensities
        
        # 发布处理后的点云消息
        pcl_msg = self.create_pointcloud2_msg(msg.header, filtered_points)
        self.pcl_pub.publish(pcl_msg)
        
        # 更新可视化
        with self.point_cloud_lock:
            points_np = filtered_points[:, :3]
            
            if self.is_first_frame:
                x_min, y_min, z_min = np.min(points_np, axis=0)
                x_max, y_max, z_max = np.max(points_np, axis=0)
                self.get_logger().info(
                    f"坐标范围: x=[{x_min:.2f}, {x_max:.2f}], "
                    f"y=[{y_min:.2f}, {y_max:.2f}], "
                    f"z=[{z_min:.2f}, {z_max:.2f}]"
                )
                
                scale = 1.0
                max_coord = max(abs(x_max), abs(x_min), abs(y_max), abs(y_min), abs(z_max), abs(z_min))
                if max_coord > 1000:
                    scale = 0.001
                    self.get_logger().info(f"应用缩放因子: {scale}")
                    points_np *= scale
            
            self.pcd.points = o3d.utility.Vector3dVector(points_np)
            self.pcd.colors = o3d.utility.Vector3dVector(colors)
            
            if self.save_enabled:
                self.save_current_frame()
            
            self.is_first_frame = False
        
        # 记录处理信息
        self.frame_count += 1
        process_time = (time.time() - start_time) * 1000
        self.get_logger().info(
            f"帧 {self.frame_count}: 原始{len(msg.points)}点 → 过滤后{len(filtered_points)}点 → "
            f"处理时间{process_time:.2f}ms"
        )

    def hsv_to_rgb(self, h, s, v):
        """HSV转RGB颜色空间"""
        h6 = h * 6
        i = np.floor(h6)
        f = h6 - i
        p = v * (1 - s)
        q = v * (1 - s * f)
        t = v * (1 - s * (1 - f))
        
        i = int(i) % 6
        if i == 0:
            return [v, t, p]
        elif i == 1:
            return [q, v, p]
        elif i == 2:
            return [p, v, t]
        elif i == 3:
            return [p, q, v]
        elif i == 4:
            return [t, p, v]
        else:
            return [v, p, q]

    def create_pointcloud2_msg(self, header, points_array):
        """创建ROS标准PointCloud2消息"""
        pcl_msg = PointCloud2()
        pcl_msg.header = header
        pcl_msg.height = 1
        pcl_msg.width = len(points_array)
        pcl_msg.fields = FIELDS
        pcl_msg.is_bigendian = False
        pcl_msg.point_step = POINT_STEP
        pcl_msg.row_step = POINT_STEP * pcl_msg.width
        pcl_msg.is_dense = False
        
        data = []
        for point in points_array:
            x, y, z, intensity = point
            data.append(struct.pack('fffB', x, y, z, int(intensity)))
        pcl_msg.data = b''.join(data)
        
        return pcl_msg

    def visualization_loop(self):
        """可视化线程主循环"""
        is_first_visualization = True
        
        while self.is_visualizing:
            with self.point_cloud_lock:
                if is_first_visualization and len(self.pcd.points) > 0:
                    self.vis.add_geometry(self.pcd)
                    
                    ctr = self.vis.get_view_control()
                    ctr.set_front([0, 0, 1])
                    ctr.set_up([0, 1, 0])
                    ctr.set_lookat([0, 0, 0])
                    ctr.set_zoom(0.4)
                    
                    is_first_visualization = False
                elif not is_first_visualization:
                    self.vis.update_geometry(self.pcd)
                    self.vis.poll_events()
                    self.vis.update_renderer()
            
            time.sleep(0.03)

    def save_current_frame(self):
        """保存当前帧点云"""
        filename = f"{self.save_path}/pointcloud_{self.frame_count:06d}.pcd"
        o3d.io.write_point_cloud(filename, self.pcd)
        self.get_logger().info(f"已保存点云: {filename}")

    def destroy_node(self):
        """节点销毁清理"""
        self.is_visualizing = False
        self.vis.destroy_window()
        super().destroy_node()


def main(args=None):
    """程序主入口"""
    rclpy.init(args=args)
    node = LidarProcessor()
    
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
